import numpy
import numpy as np
import matplotlib.pyplot as plt
from pandas import read_csv
import math
import pandas as pd
from keras.models import Sequential
from keras.layers import Dense,Dropout
from keras.callbacks import EarlyStopping
from keras.layers import LSTM
from keras.constraints import maxnorm
from keras.regularizers import l1_l2,l1,l2
from keras.layers import RNN, SimpleRNN, SimpleRNNCell
from sklearn.preprocessing import MinMaxScaler
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
from keras.utils.vis_utils import plot_model



def mean_absolute_percentage_error(y_true, y_pred):
    # 平均绝对百分比误差（MAPE）的计算
    y_true, y_pred = np.array(y_true), np.array(y_pred)
    return np.mean(np.abs((y_true - y_pred) / y_true)) * 100

if __name__ == '__main__':
    # 加载数据
    dataframe = read_csv('E:\lyf_ML_Drought\coding\ML_Drought_Prediction\indices_caculate\\result\ROW_SPEI\ROW_SPEI-12\SPEI-12_52533.txt', header=None, names=('TIME','SPEI-1'))
    dataframe = dataframe.set_index(['TIME'], drop=True) # 把日期作为索引

    # 创建训练和测试数据集
    # train, test = create_dataset(np.array(dataframe), 0.9)
    train = dataframe.iloc[:int(0.9*len(dataframe))]
    test = dataframe.iloc[int(0.9*len(dataframe)):]

    # 归一化
    scaler = MinMaxScaler(feature_range=(-1, 1))
    train_sc = scaler.fit_transform(train.values.reshape(-1, 1))
    test_sc = scaler.transform(test.values.reshape(-1, 1))


    # 在构造LSTM时，我们将使用pandas shift函数将整列往后移动1个位置.以便得到训练数据和标签数据.然后我们需要将所有训练数据和标签数据转换成3维张量
    train_sc_df = pd.DataFrame(train_sc, columns=['Y'], index=train.index)
    test_sc_df = pd.DataFrame(test_sc, columns=['Y'], index=test.index)

    for s in range(1, 2):
        train_sc_df['X_{}'.format(s)] = train_sc_df['Y'].shift(s)
        test_sc_df['X_{}'.format(s)] = test_sc_df['Y'].shift(s)

    X_train = train_sc_df.dropna().drop('Y', axis=1)
    y_train = train_sc_df.dropna().drop('X_1', axis=1)

    X_test = test_sc_df.dropna().drop('Y', axis=1)
    y_test = test_sc_df.dropna().drop('X_1', axis=1)

    X_train = X_train.values
    y_train = y_train.values

    X_test = X_test.values
    y_test = y_test.values

    X_train_lmse = X_train.reshape(X_train.shape[0], X_train.shape[1], 1)
    X_test_lmse = X_test.reshape(X_test.shape[0], X_test.shape[1], 1)

    print('Train shape: ', X_train_lmse.shape)
    print('Test shape: ', X_test_lmse.shape)

    """
    LSTM有1个输入层
    1个带有7个LSTM神经元的隐藏层
    一个进行单值预测的输出层
    LSTM神经元的激活函数为tanh
    设置早期停止当损失值不再改善时
    模型训练100个周期,bach_size=1
    """

    lstm_model = Sequential()
    # lstm_model.add(
    #     LSTM(7, input_shape=(1, X_train_lmse.shape[1]), kernel_initializer='lecun_uniform',
    #          return_sequences=False,
    #          bias_regularizer=l1_l2(l1=0.01, l2=0.01),
    #          kernel_constraint=maxnorm(5), kernel_regularizer=l1_l2(l1=0.01, l2=0.01),))
    lstm_model.add(Dense(48, activation='tanh'))
    # lstm_model.add(Dropout(0.2)) # Dropout是遗忘概率
    lstm_model.add(Dense(1))
    lstm_model.compile(loss='mean_squared_error', optimizer='adam')
    early_stop = EarlyStopping(monitor='loss', patience=5, verbose=1)
    history_lstm_model = lstm_model.fit(X_train_lmse, y_train, epochs=3000, batch_size=50, verbose=2, shuffle=False,
                                        callbacks=[early_stop])

    y_pred_test_lstm = lstm_model.predict(X_test_lmse)
    y_train_pred_lstm = lstm_model.predict(X_train_lmse)

    y_train_true_lstm = scaler.inverse_transform(y_train)
    y_pred_true_lstm = scaler.inverse_transform(y_test)
    y_train_lstm = scaler.inverse_transform(y_train_pred_lstm)
    y_pred_lstm = scaler.inverse_transform(y_pred_test_lstm)

    print("train rmse: ", numpy.sqrt(mean_squared_error(y_train_true_lstm, y_train_lstm)))
    print("train mae: ", mean_absolute_error(y_train_true_lstm, y_train_lstm))
    print("train mape: ", mean_absolute_percentage_error(y_train_true_lstm, y_train_lstm))
    print("train r2: ", r2_score(y_train_true_lstm, y_train_lstm))

    print("predict rmse: ", numpy.sqrt(mean_squared_error(y_pred_true_lstm, y_pred_lstm)))
    print("predict mae: ", mean_absolute_error(y_pred_true_lstm, y_pred_lstm))
    print("predict mape: ", mean_absolute_percentage_error(y_pred_true_lstm, y_pred_lstm))
    print("predict r2: ", r2_score(y_pred_true_lstm, y_pred_lstm))

    plt.figure()
    plt.plot(y_pred_true_lstm, label='True',c="black")
    plt.plot(y_pred_lstm, label='LSTM',c="red")
    plt.title("LSTM's Prediction")
    plt.xlabel('Month')
    plt.ylabel('SPEI-1')
    plt.legend()
    plt.show();
